#!/bin/bash
# date 2023-06-15
# used for erasing disk for etcd and rook.
# code 0: check OK
# code 1: invalid nums of arguments. 1 and none.
# code 2: invalid arguments. arguments must be a disk.
# code 3: fio exec failed.
# code 4: parse fio-benchmark failed.
# code 5: fio-benchmark greater than limit_ns.
# code 6: exec dd failed.
# code 7: small files IOPS smaller than limit_rate.
#set -x

dd_save_dir="/etc/benchmark-dd"
dd_limit_rate=4  # 4MB/s
benchmark_save_dir="/etc/benchmark-fio"
limit_ns=10000000
CURRENT_VALUE=0
small_file_level=""

function log_info(){
  echo -e "{\"code\":$1, \"data\":{\"fio\":{\"CURRENT_VALUE\":$2, \"LIMIT_VALUE\":${limit_ns}}, \"dd\":{\"CURRENT_VALUE\":\"$3\", \"LIMIT_VALUE\":\"4 MB/s\"}}, \"msg\":\"none\"}"
}

function log_warn(){
  echo -e "{\"code\":$1, \"data\":{\"fio\":{\"CURRENT_VALUE\":$2, \"LIMIT_VALUE\":${limit_ns}}, \"dd\":{\"CURRENT_VALUE\":\"$3\", \"LIMIT_VALUE\":\"4 MB/s\"}}, \"msg\":\"$4\"}"
}

function log_error(){
  echo -e "{\"code\":$1, \"data\":{}, \"msg\":\"$2\"}"
}

function fio_check()
{
  name=($1)
  fio_check=$(fio --rw=write --ioengine=sync --fdatasync=1 --size=22m --bs=2300 --name=${name} --output-format=json --output=${benchmark_save_dir})
#  fio_check=$(fio --rw=write --ioengine=sync --fdatasync=1 --size=22m --bs=2300 --directory=${benchmark_test_dir} --name=check --output-format=json --output=${benchmark_save_dir})
  if [ $? -ne 0 ]; then
    log_error 3 "fio exec failed"
    exit 3
  fi
  if [ -e ${benchmark_save_dir} ];then
    CURRENT_VALUE=$(cat ${benchmark_save_dir} | jq -r '.jobs[].sync.lat_ns.percentile."99.000000"')
    if [ $? -ne 0 ]; then
      CURRENT_VALUE=$(grep 99.00000 ${benchmark_save_dir} | tail -1 | awk -F"[:,]" '{print $2}')
      if [ $? -ne 0 ]; then
        log_error 4 "parse ${benchmark_save_dir} failed"
        exit 4
      fi
    fi
    if [ ${CURRENT_VALUE} -gt ${limit_ns} ]; then
      log_warn 5 ${CURRENT_VALUE} "TBD" "sync I/O latency too high"
      exit 5
    fi
  else
    log_error 3 "${benchmark_save_dir} NOT FOUND"
    exit 3
  fi
}

function check_sync_bench() {
  if [ -z ${dev} ];then
    cd /
    fio_check "fio-tempfile"
  else
    fdisk -l ${dev} 1>/dev/null 2>&1
    if [ $? -ne 0 ]; then
      log_error 2 "invalid device"
      exit 2
    fi
    benchmark_save_dir="$benchmark_save_dir-$(echo $dev| cut -d/ -f3)"
    fio_check ${dev}
  fi
}

function evaluate_level()
{
    local bs=$1
    local count=$2
    local excellent_secs=$3
    local device=$4
    local level=0
    local timeout_secs=$((${excellent_secs} * 4))   # 注意，这里取值4倍

    res=$(timeout ${timeout_secs}s dd if=/dev/zero of=${device} oflag=direct oflag=sync bs=${bs} count=${count} 2>&1)
    if [ $? -ne 0 ]; then
        echo $res > ${dd_save_dir}
        echo -e "disk ${device} IOPS check \nfinished in $(date)" >> ${dd_save_dir}
        if [ ${device} = "/root/for_benchmarking_io"];then
          rm -f ${device}
        fi
        log_error 6 "dd exec failed"
        exit 6
    fi
    echo ${res} > ${dd_save_dir}
    echo -e "disk ${device} IOPS check \nfinished in $(date)" >> ${dd_save_dir}
    if [ ${device} = "/root/for_benchmarking_io" ];then
      rm -f ${device}
    fi
    #200+0 records in 200+0 records out 819200 bytes (819 kB, 800 KiB) copied, 0.0795682 s, 10.3 MB/s
    small_file_level=$(echo $res | grep -o '[0-9.]*..B/.$')
    if [[ ${small_file_level} =~ 'MB' ]];then
      dd_limit_rate=4
    elif [[ ${small_file_level} =~ 'kB' ]] || [[ ${small_file_level} =~ 'KB' ]];then
      dd_limit_rate=4000
    elif [[ ${small_file_level} =~ 'GB' ]];then
      dd_limit_rate=0.004
    else
      log_error 6 "unknown dd value"
      exit 6
    fi

    level_value=$(echo ${small_file_level}| grep -o '[0-9.]*')
    if [ $(echo "${level_value} >= ${dd_limit_rate}"|bc) -eq 0 ];then
      log_warn 7 ${CURRENT_VALUE} "${small_file_level}" "small files IOPS check failed"
      exit 7
    fi
}

function check_IOPS_bench() {
  if [ -z ${dev} ];then
    device="/root/for_benchmarking_io"
  else
    fdisk -l ${dev} 1>/dev/null 2>&1
    if [ $? -ne 0 ]; then
      log_error 2 "invalid device"
      exit 2
    fi
    device=${dev}
    dd_save_dir="${dd_save_dir}-$(echo $dev| cut -d/ -f3)"
  fi
  SMALL_FILE_TEST_BS=4K
  SMALL_FILE_TEST_COUNT=2000
  SMALL_FILE_TEST_EXCELLENT_SECS=10
  evaluate_level ${SMALL_FILE_TEST_BS} ${SMALL_FILE_TEST_COUNT} ${SMALL_FILE_TEST_EXCELLENT_SECS} ${device}
}

function main() {
    check_sync_bench
    check_IOPS_bench
    log_info 0 ${CURRENT_VALUE} "${small_file_level}"
}

if [ $# -gt 1 ]; then
  log_error 1 "invalid nums of arguments. just 0 or 1"
  exit 1
fi
dev=($1)
main
exit 0
